import os
import glob
import argparse
from gmssl import *

def test_sm4_cbc(key_file, iv_file, plain_file, cipher_file):
    key = read_binary_file_(key_file) 
    iv = read_binary_file_(iv_file) 
    plaintext = read_binary_file_(plain_file) 
    ciphertext_hex = read_binary_file_(cipher_file) 
    
    sm4_cbc = Sm4Cbc(key, iv, DO_ENCRYPT)
    ciphertext = sm4_cbc.update(plaintext)
    ciphertext += sm4_cbc.finish()
    cipher_std = ciphertext[0:len(plaintext)]
    #print("plaintext", plaintext.hex())
    #print("ciphertext", cipher_std.hex())
    assert ciphertext_hex == cipher_std, "加密失败"

def test_sm4_ctr (key_file, iv_file, plain_file, cipher_file):
    key = read_binary_file_(key_file) 
    iv = read_binary_file_(iv_file) 
    plaintext = read_binary_file_(plain_file) 
    ciphertext_hex = read_binary_file_(cipher_file) 

    sm4_enc = Sm4Ctr(key, iv)
    ciphertext = sm4_enc.update(plaintext)
    ciphertext += sm4_enc.finish()
    
    #print("plaintext", plaintext.hex())
    #print("ciphertext_hex", ciphertext_hex.hex())
    #print("ciphertext", ciphertext.hex())
    assert ciphertext_hex == ciphertext, "加密失败"

def read_binary_file_(file_path):
    with open(file_path, 'rb') as f:
        binary_data = f.read()
        return binary_data


def read_binary_file(file_path):
    """读取二进制文件内容并转换为字符串"""
    with open(file_path, 'rb') as f:
        binary_data = f.read()
        # 尝试解码为字符串，如果失败则用十六进制表示
        return binary_data.hex().upper()

def main():
    parser = argparse.ArgumentParser(description='遍历目录中的二进制文件')
    parser.add_argument('directory', help='目标目录路径')
    parser.add_argument('mode', choices=['enc', 'dec'], default='enc', help='运行模式')
    args = parser.parse_args()
    directory = args.directory
    mode = args.mode
    print (f"模式:{mode}\n")

    alg_mode = ['CBC', 'CFB', 'CTR', 'ECB', 'OFB']
    cipher_mode = ['cipher_cbc', 'cipher_cfb', 'cipher_ctr', 'cipher_ecb', 'cipher_ofb']
    replain_mode = ['replain_cbc', 'replain_cfb', 'replain_ctr', 'replain_ecb', 'replain_ofb']
    enc_output_flie = ['sm4_cbc加密标准数据.txt', 'sm4_cfb加密标准数据.txt','sm4_ctr加密标准数据.txt','sm4_ecb加密标准数据.txt','sm4_ofb加密标准数据.txt']
    dec_output_flie = ['sm4_cbc解密标准数据.txt', 'sm4_cfb解密标准数据.txt','sm4_ctr解密标准数据.txt','sm4_ecb解密标准数据.txt','sm4_ofb解密标准数据.txt']

    if not os.path.isdir(directory):
        print(f"错误：目录不存在 - {directory}")
        return
    
    # 创建输出内容
    merged_content = ""
    for j in range(5):
        
        iv_file = os.path.join(directory, f'iv')
        key_file = os.path.join(directory, f'key')
        cipher_file = os.path.join(directory, f'{cipher_mode[j]}')
        if mode == 'enc' :  #加密
            plain_file = os.path.join(directory, f'plain')
        else:
            plain_file = os.path.join(directory, f'{replain_mode[j]}')
        match j:
                case 0:
                    test_sm4_cbc(key_file, iv_file, plain_file, cipher_file)
                case 2:
                    test_sm4_ctr(key_file, iv_file, plain_file, cipher_file)
        # 读取文件内容
        if j != 3 : 
            iv_content = read_binary_file(iv_file)
        key_content = read_binary_file(key_file)
        cipher_content = read_binary_file(cipher_file)
        if mode == 'enc':
            plain_content = read_binary_file(plain_file)
            length = int(len(plain_content) / 2)
        else:
            replain_content = read_binary_file(plain_file)
            length = int(len(replain_content) / 2)

        # 添加到合并内容
        merged_content += f"{alg_mode[j]}\n"
        merged_content += f"密钥=\n{key_content}\n"
        if j != 3 : 
            merged_content += f"IV=\n{iv_content}\n"
        merged_content += f"明文长度= {length:08x}\n"
        if mode == 'enc':
            merged_content += f"明文=\n{plain_content}\n"
            merged_content += f"密文=\n{cipher_content}\n"
        else:
            merged_content += f"密文=\n{cipher_content}\n"
            merged_content += f"明文=\n{replain_content}\n"
        merged_content += "\n"
        
    
    # 写入到txt文件
    if mode == 'enc':
        output_file = "sm4加密速率.txt" 
    else:
        output_file = "sm4解密速率.txt" 


    with open(output_file, 'w', encoding='utf-8') as f:
        f.write(merged_content)
        
    print(f"合并完成，结果已保存到 {output_file}")

if __name__ == "__main__":
    main()
